"""
五子棋 DQN 系统 - 增强版主程序
将训练模式和游戏模式分离，提高代码可读性
"""

import argparse
from gomoku_trainer import GomokuTrainer
from gomoku_gui import GomokuGame


def main():
    """主函数：选择训练模式或游戏模式"""
    parser = argparse.ArgumentParser(description='五子棋 DQN 系统 - 增强版')
    parser.add_argument('--mode', choices=['train', 'play'], default='play',
                        help='选择模式: train=训练AI, play=与AI对弈 (默认: play)')
    parser.add_argument('--board_size', type=int, default=11,
                        help='棋盘大小 (默认: 11)')
    parser.add_argument('--episodes', type=int, default=1000,
                        help='训练回合数 (仅训练模式, 默认: 1000)')
    parser.add_argument('--model_path', type=str, default=None,
                        help='AI模型路径 (仅游戏模式)')
    parser.add_argument('--resume', type=str, default=None,
                        help='从指定检查点恢复训练')

    args = parser.parse_args()

    if args.mode == 'train':
        print("🔥 开始增强版训练模式...")
        trainer = GomokuTrainer(board_size=args.board_size)
        trainer.train(num_episodes=args.episodes, resume_model_path=args.resume)
        print("增强版训练完成！")

    elif args.mode == 'play':
        print("开始游戏模式...")
        if args.model_path:
            print(f"加载AI模型: {args.model_path}")
        else:
            print("未指定模型路径，将进行双人对弈模式")

        game = GomokuGame(
            model_path=args.model_path,
            board_size=args.board_size,
            ai_player_num=-1
        )

        if game.gui_enabled:
            game.play()
        else:
            print("GUI未启用，无法开始游戏")


if __name__ == "__main__":
    main()